package parser;

import java.util.Arrays;
import java.util.List;

import se.lth.cs.srl.CompletePipeline;
import se.lth.cs.srl.corpus.Predicate;
import se.lth.cs.srl.corpus.Sentence;

public class LTHSRLAnnotation extends ParserAnnotation {

	private CompletePipeline parser;
	private String[] lemma;
	private String[] tokens;
	private String[] pos;
	private Integer[] heads;
	private String[] deprels;
	private String[] srl;

	/**
	 * This parser does NOT do sentence detection.
	 */
	public LTHSRLAnnotation(CompletePipeline parser, String doc) {
		super(doc);
		this.parser = parser;
	}

	@Override
	protected void parseImplementation() throws Exception {
		Sentence sent = parser.parse(document);
		tokens = sent.getFormArray();
		tokens = Arrays.copyOfRange(tokens, 1, tokens.length);
		pos = sent.getPOSArray();
		pos = Arrays.copyOfRange(pos, 1, pos.length);
		int size = sent.size() - 1;
		lemma = new String[size];
		heads = new Integer[size];
		srl = new String[size];
		deprels = new String[size];
		for (int i = 0; i < lemma.length; i++) {
			lemma[i] = sent.get(i + 1).getLemma();
			if (lemma[i] == null)
				lemma[i] = tokens[i];
			heads[i] = sent.get(i + 1).getHeadId();
			deprels[i] = sent.get(i + 1).getDeprel();
			if (sent.get(i+1) instanceof Predicate)
				srl[i] = "P";
			else {
				for (Predicate p : sent.getPredicates()) {
					if (p.getArgumentTag(sent.get(i+1)) != null)
						srl[i] = p.getArgumentTag(sent.get(i+1));
				}
			}
			if (srl[i] == null)
				srl[i] = "-";
		}
	}

	@Override
	public List<String> getSentences() {
		return Arrays.asList(document);
	}

	@Override
	public int getNbrSentences() {
		return 1;
	}

	@Override
	public List<String> getTokens(int sent) {
		return Arrays.asList(tokens);
	}

	@Override
	public List<Integer> getTokenIndexes(int sent) {
		throw new Error("Token indices is not implemented for LTHSRL parser.");
	}

	@Override
	public List<String> getPOS(int sent) {
		return Arrays.asList(pos);
	}

	@Override
	public List<String> getLemmas(int sent) {
		return Arrays.asList(lemma);
	}

	public List<String> getDeprel(int sent) {
		return Arrays.asList(deprels);
	}

	public List<String> getSRL(int sent) {
		return Arrays.asList(srl);
	}

	public List<Integer> getHeadIds(int sent) {
		return Arrays.asList(heads);
	}
}
